package com.liuhengjia.langchain4japply.config;

import com.liuhengjia.langchain4japply.config.assistant.Assistant;
import com.liuhengjia.langchain4japply.config.assistant.MemoryIsolationAssistant;
import com.liuhengjia.langchain4japply.config.assistant.StructuredOutputAssistant;
import com.liuhengjia.langchain4japply.tools.CalculatorPlus;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import dev.langchain4j.web.search.WebSearchTool;
import dev.langchain4j.web.search.searchapi.SearchApiWebSearchEngine;
import lombok.RequiredArgsConstructor;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

/**
 * 聊天助手 Bean
 */
@Configuration
@RequiredArgsConstructor
public class AssistantConfig {
    private final ChatLanguageModel chatLanguageModel;

    /**
     * 基础助手 Bean
     *
     * @return baseAssistant
     */
    @Bean
    public Assistant baseAssistant() {
        return AiServices.builder(Assistant.class)
                .chatLanguageModel(chatLanguageModel)
                .build();
    }

    /**
     * 会话记忆助手 Bean
     *
     * @return memoryAssistant
     */
    @Bean
    public Assistant memoryAssistant() {
        return AiServices.builder(Assistant.class)
                .chatLanguageModel(chatLanguageModel)
                .chatMemory(MessageWindowChatMemory.builder().maxMessages(10).build())
                .build();
    }

    /**
     * 会话隔离助手 Bean
     *
     * @return memoryIsolationAssistant
     */
    @Bean
    public Assistant memoryIsolationAssistant() {
        return AiServices.builder(MemoryIsolationAssistant.class)
                .chatLanguageModel(chatLanguageModel)
                .chatMemoryProvider(memoryId -> MessageWindowChatMemory.builder().maxMessages(10).build())
                .build();
    }

    /**
     * LangChain4j 提供的基于内存的数据库
     *
     * @return EmbeddingStore<TextSegment>
     */
    @Bean
    public EmbeddingStore<TextSegment> embeddingStore() {
        return new InMemoryEmbeddingStore<>();
    }

    /**
     * RAG 助手 Bean
     *
     * @return ragAssistant
     */
    @Bean
//    public Assistant ragAssistant(EmbeddingStore<TextSegment> embeddingStore) {
    public Assistant ragAssistant(EmbeddingStore<TextSegment> pgVectorEmbeddingStore) {
        return AiServices.builder(MemoryIsolationAssistant.class)
                .chatLanguageModel(chatLanguageModel)
                .chatMemoryProvider(memoryId -> MessageWindowChatMemory.builder().maxMessages(10).build())
//                .contentRetriever(EmbeddingStoreContentRetriever.from(embeddingStore))// 内置
                .contentRetriever(EmbeddingStoreContentRetriever.from(pgVectorEmbeddingStore))// PgVector
                .build();
    }

    /**
     * Function 助手 Bean
     *
     * @return functionAssistant
     */
    @Bean
    public Assistant functionAssistant() {
        return AiServices.builder(Assistant.class)
                .chatLanguageModel(chatLanguageModel)
                .tools(new CalculatorPlus())
                .build();
    }

    /**
     * Web Search 助手 Bean
     *
     * @return webSearchAssistant
     */
    @Bean
    public Assistant webSearchAssistant(SearchApiWebSearchEngine searchApiWebSearchEngine) {
        return AiServices.builder(Assistant.class)
                .chatLanguageModel(chatLanguageModel)
                .tools(new WebSearchTool(searchApiWebSearchEngine))
                .build();
    }

    /**
     * Web Search 助手 Bean
     *
     * @return webSearchAssistant
     */
    @Bean
    public Assistant structuredOutputAssistant() {
        return AiServices.builder(StructuredOutputAssistant.class)
                .chatLanguageModel(chatLanguageModel)
                .build();
    }
}
